#!/usr/bin/env python
# coding: utf-8

# In[ ]:


# 导入相关包
from utils import DGraphFin
from utils.utils import prepare_folder
from utils.evaluator import Evaluator
import pandas as pd

import torch
import torch.nn.functional as F
import torch.nn as nn

import torch_geometric.transforms as T

import numpy as np
from torch_geometric.data import Data
import os

# 随机种子
torch.manual_seed(666)

#设置gpu设备
device = 0
device = f'cuda:{device}' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)

# CUDA is available
print(torch.cuda.is_available())


# In[ ]:


# 定义网络模型

class MLP(torch.nn.Module):
    def __init__(self
                 , in_channels
                 , hidden_channels
                 , out_channels
                 , num_layers
                 , dropout
                 , batchnorm=True):
        super(MLP, self).__init__()
        self.lins = torch.nn.ModuleList()
        self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
        self.batchnorm = batchnorm
        if self.batchnorm:
            self.bns = torch.nn.ModuleList()
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
            if self.batchnorm:
                self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.lins.append(torch.nn.Linear(hidden_channels, out_channels))

        self.dropout = dropout

    def reset_parameters(self):
        for lin in self.lins:
            lin.reset_parameters()
        if self.batchnorm:
            for bn in self.bns:
                bn.reset_parameters()

    def forward(self, x):    
        for i, lin in enumerate(self.lins[:-1]):
            x = lin(x)
            if self.batchnorm:
                x = self.bns[i](x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.lins[-1](x)
        return F.log_softmax(x, dim=-1)

def train(model, data, train_idx, optimizer):
     # data.y is labels of shape (N, ) 
    model.train()

    optimizer.zero_grad()
    
    out = model(data.x[train_idx])

    loss = F.nll_loss(out, data.y[train_idx])
    loss.backward()
    optimizer.step()

    return loss.item()

def test(model, data, split_idx, evaluator):
    # data.y is labels of shape (N, )
    with torch.no_grad():
        model.eval()

        losses, eval_results = dict(), dict()
        for key in ['train', 'valid']:
            node_id = split_idx[key]
            
            out = model(data.x[node_id])
            y_pred = out.exp()  # (N,num_classes)
            
            losses[key] = F.nll_loss(out, data.y[node_id]).item()
            eval_results[key] = evaluator.eval(data.y[node_id], y_pred)[eval_metric]

    return eval_results, losses, y_pred

def predict(data,node_id):
    """
    加载模型和模型预测
    :param node_id: int, 需要进行预测节点的下标
    :return: tensor, 类0以及类1的概率, torch.size[1,2]
    """
    model = MLP(in_channels = 20, hidden_channels = 128, out_channels = 2, num_layers=2, dropout = 0, batchnorm = False)
    model.load_state_dict(torch.load(save_dir+'/model_mlp.pt')) #载入验证集上表现最好的模型
    
    # -------------------------- 实现模型预测部分的代码 ---------------------------
    with torch.no_grad():
        model.eval()
        out = model(data.x[node_id])
        y_pred = out.exp()  # (N,num_classes)
        
    return y_pred


# In[ ]:


path='./datasets/632d74d4e2843a53167ee9a1-momodel/' #数据保存路径
save_dir='./results/' #模型保存路径
dataset_name='DGraph'
dataset = DGraphFin(root=path, name=dataset_name, transform=T.ToSparseTensor())

nlabels = dataset.num_classes
if dataset_name in ['DGraph']:
    nlabels = 2    #本实验中仅需预测类0和类1

data = dataset[0]
data.adj_t = data.adj_t.to_symmetric() #将有向图转化为无向图


if dataset_name in ['DGraph']:
    x = data.x
    x = (x - x.mean(0)) / x.std(0)
    data.x = x
if data.y.dim() == 2:
    data.y = data.y.squeeze(1)

split_idx = {'train': data.train_mask, 'valid': data.valid_mask, 'test': data.test_mask}  #划分训练集，验证集

train_idx = split_idx['train']
result_dir = prepare_folder(dataset_name,'mlp')

# 查看数据维度
print(data)
print(data.x.shape)  #feature
print(data.y.shape)  #label


# In[ ]:


# 定义网络模型
model = MLP(in_channels=data.x.size(-1), hidden_channels=128, out_channels=2, num_layers=2, dropout=0.0, batchnorm=False)
print(f'Model MLP initialized')
eval_metric = 'auc'  #使用AUC衡量指标
evaluator = Evaluator(eval_metric)
epochs = 200


# In[ ]:


# 训练网络模型
print(sum(p.numel() for p in model.parameters()))  #模型总参数量

model.reset_parameters()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-7)
best_valid = 0
min_valid_loss = 1e8

Epochs = []
Loss   = []
Train_AUC = []
Valid_AUC = []

for epoch in range(1,epochs + 1):
    loss = train(model, data, train_idx, optimizer)
    eval_results, losses, out = test(model, data, split_idx, evaluator)
    train_eval, valid_eval = eval_results['train'], eval_results['valid']
    train_loss, valid_loss = losses['train'], losses['valid']

    if valid_loss < min_valid_loss:
        min_valid_loss = valid_loss
        torch.save(model.state_dict(), save_dir+'/model_mlp.pt') #将表现最好的模型保存

    print(f'Epoch: {epoch:02d}, '
            f'Loss: {loss:.4f}, '
            f'Train: {100 * train_eval:.3f}, ' # 我们将AUC值乘上100，使其在0-100的区间内
            f'Valid: {100 * valid_eval:.3f} ')

    print('{{"metric": "Loss", "value": {:.4f}, "epoch": {} }}'.format(loss,epoch))
    print('{{"metric": "Train AUC", "value": {:.4f}, "epoch": {} }}'.format(100*train_eval,epoch))
    print('{{"metric": "Valid AUC", "value": {:.4f}, "epoch": {} }}'.format(100*valid_eval,epoch))

    Epochs.append(epoch)
    Loss.append(loss)
    Train_AUC.append(train_eval)
    Valid_AUC.append(valid_eval)


# In[ ]:


dataLogger = pd.DataFrame()
dataLogger.insert(0,"Epochs", Epochs)
dataLogger.insert(1,"Loss",Loss)
dataLogger.insert(2,"Train_AUC",Train_AUC)
dataLogger.insert(2,"Valid_AUC",Valid_AUC)
dataLogger.to_excel('MLP_trainingLog.xlsx',float_format="%.4f",index=False)

